{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "91053b80-86fb-4736-a96e-a6062a36c078",
   "metadata": {},
   "source": [
    "# 千帆 与 Weights & Biases"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "627e4d20-6df5-4191-94ae-34b4524ce9e5",
   "metadata": {},
   "source": [
    "Weights & Biases 提供了一个平台，可以对大模型的调用过程进行可视化的展现，跟踪每次请求对资源的消耗等等，尤其是在开发 Agent 等涉及到复杂的模型调用时，可以帮助开发者更好地观察模型效果并进一步去改进模型。\n",
    "\n",
    "本文将介绍如何利用 LangChain，将千帆SDK的调用导入 Weights & Biases。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce0f4e5f-896a-4d92-910d-724fd25f993b",
   "metadata": {},
   "source": [
    "## 1. 初始化"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0b3578de-3103-4846-8fd5-e2331d2efbf2",
   "metadata": {},
   "source": [
    "为了能够使用 Weights & Biases，需要先安装 wandb 和 langchain，以及qianfan：对应的版本如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19434a96",
   "metadata": {
    "vscode": {
     "languageId": "bat"
    }
   },
   "outputs": [],
   "source": [
    "langchain                     0.0.335\n",
    "qianfan                       0.1.3\n",
    "wandb                         0.16.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeff9bc6-043d-4a41-a543-4052556e1b95",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install qianfan\n",
    "!pip install wandb\n",
    "!pip install langchain"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "62215460-8197-4ebb-b608-9266e7385034",
   "metadata": {},
   "source": [
    "之后在代码中我们需要开启 LangChain 对 WanDB 的 Tracing 功能。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9a043ffe-daf3-4fb4-b64c-3bbf8d7b729f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import wandb\n",
    "import qianfan\n",
    "\n",
    "# 初始化千帆 SDK\n",
    "os.environ[\"QIANFAN_AK\"] = \"your_ak\"\n",
    "os.environ[\"QIANFAN_SK\"] = \"your_sk\"\n",
    "# 开启 Langchain 对 WanDB 的 Tracing 功能\n",
    "os.environ[\"LANGCHAIN_WANDB_TRACING\"] = \"true\"\n",
    "# 设置 WanDB 的项目名称\n",
    "os.environ[\"WANDB_PROJECT\"] = \"my-wandb-llm-project\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39123465-8b48-42f7-848d-2e55514eee17",
   "metadata": {},
   "source": [
    "## 2. 使用"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25a3b97a-dca1-4579-9b67-aebaf7658a68",
   "metadata": {},
   "source": [
    "使用方法与正常使用 LangChain 一致，唯一的区别是在调用模型时，需要用 wandb 进行包装。\n",
    "\n",
    "在第一次使用时，需要填入 wandb 的 api key，请遵循输出的提示进行注册、获取 key 等操作。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f1f9402-364c-4a02-ad78-44a2071592de",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import QianfanChatEndpoint\n",
    "\n",
    "llm = QianfanChatEndpoint(model=\"ERNIE-Bot\")\n",
    "\n",
    "with wandb.init(project='langchain-tracing') as run:\n",
    "    resp = llm.invoke(\"hello\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0b6f749-dfca-46aa-bace-32c71d657ee4",
   "metadata": {},
   "source": [
    "运行完成后，会输出该次运行的链接，可以在 wandb 的控制台上查看到 trace 信息。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd2fd27c-2c5a-455b-8081-234e27689bd4",
   "metadata": {},
   "source": [
    "## 3. 跟踪 Agent"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "efc22231-23f9-4ae3-ab29-6505b090aad6",
   "metadata": {},
   "source": [
    "这里为了展示一个更为复杂的用例，我们先创建一个 Agent，采用的是在 [这一 cookbook](https://github.com/baidubce/bce-qianfan-sdk/blob/main/cookbook/qianfan_agent/qianfan_agent.ipynb) 中所构建的利用 function call 实现的获取天气信息的 Agent，在此不再重复介绍。\n",
    "\n",
    "这边也可以替换能任何通过 LangChain 实现的组件。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a2f35645-44a1-4945-abe9-8a5568fd9602",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import QianfanChatEndpoint\n",
    "from langchain.agents import load_tools, initialize_agent, AgentType\n",
    "from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain.agents import tool\n",
    "from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser\n",
    "from langchain.agents.format_scratchpad import format_to_openai_functions\n",
    "\n",
    "llm = QianfanChatEndpoint(model=\"ERNIE-Bot\") \n",
    "prompt = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", \"You are very powerful assistant, but bad at get today's temperature of location.\"),\n",
    "    (\"user\", \"{input}\"),\n",
    "    MessagesPlaceholder(variable_name=\"agent_scratchpad\")\n",
    "])\n",
    "\n",
    "FUNCTION_SCHEMA_GET_WEATHER = {\n",
    "    \"name\": \"get_current_weather\",\n",
    "    \"description\": \"获得指定地点的天气\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "        \"location\": {\n",
    "            \"type\": \"string\",\n",
    "            \"description\": \"省，市名，例如：河北省，石家庄\"\n",
    "        },\n",
    "        \"unit\": {\n",
    "            \"type\": \"string\",\n",
    "            \"enum\": [\"摄氏度\", \"华氏度\"]\n",
    "        }\n",
    "        },\n",
    "        \"required\": [\"location\"]\n",
    "    }\n",
    "}\n",
    "\n",
    "@tool\n",
    "def get_current_weather(location: str) -> str:\n",
    "    \"\"\"Returns current temperature of location.\"\"\"\n",
    "    return \"25\"\n",
    "\n",
    "tools = [get_current_weather]\n",
    "\n",
    "llm_with_tools = llm.bind(\n",
    "    functions=[\n",
    "        FUNCTION_SCHEMA_GET_WEATHER,\n",
    "    ]\n",
    ")\n",
    "p = {\n",
    "    \"input\": lambda x: x[\"input\"],\n",
    "    \"agent_scratchpad\": lambda x: format_to_openai_functions(x['intermediate_steps']),\n",
    "} \n",
    "agent = p | prompt | llm_with_tools | OpenAIFunctionsAgentOutputParser()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd5268bd-6965-418d-9bff-c200e30d9103",
   "metadata": {},
   "source": [
    "创建完 Agent 后，同样的与正常使用基本一致，只需要将需要跟踪的部分用 wandb 对象进行包裹即可。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6e0bd87-7a79-4aae-a3dc-cae7cadbbfd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.schema.agent import AgentFinish\n",
    "intermediate_steps = []\n",
    "\n",
    "with wandb.init(project='langchain-tracing') as run:\n",
    "    while True:\n",
    "        output = agent.invoke({\n",
    "            \"input\": \"上海市今天天气如何？\",\n",
    "            \"intermediate_steps\": intermediate_steps\n",
    "        })\n",
    "        \n",
    "        if isinstance(output, AgentFinish):\n",
    "            final_result = output.return_values[\"output\"]\n",
    "            break\n",
    "        else:\n",
    "            print(output.tool, output.tool_input)\n",
    "            tool = {\n",
    "                \"get_current_weather\": get_current_weather,\n",
    "            }[output.tool]\n",
    "            observation = tool.run(output.tool_input)\n",
    "            intermediate_steps.append((output, observation))\n",
    "print(final_result)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9fcb6224-e26a-4211-b185-6d6442f5d1ec",
   "metadata": {},
   "source": [
    "从输出中就可以拿到这次跟踪的链接，可以从控制台看到 Agent 运行的每一个步骤和结果，以及各类统计信息。\n",
    "\n",
    "至此，我们就完成了千帆与 wandb 的结合，可以在 wandb 的控制台上看到千帆模型的执行结果，希望能对您大模型的开发带来帮助！\n",
    "![link.img](./img/tracks.jpg)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "58f7cb64c3a06383b7f18d2a11305edccbad427293a2b4afa7abe8bfc810d4bb"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
